// dllmain.cpp : 定义 DLL 应用程序的入口点。
// 目前包含对话框函数、堆函数、以及dll主函数
#include "pch.h"
#include <cstdio>
#include <cstdlib>
#include <iostream>
#include <string>
#include <windows.h>
#include <unordered_set>
#include <detours.h>
#include "head.h"  // 头文件，包括定义的各函数的编号、info的结构

#pragma comment(lib, "detours.lib")
#pragma comment (lib, "ws2_32.lib")

using namespace std;

extern struct info;

info sendInfo;
SYSTEMTIME st;  // 系统时间
HANDLE hSemaphore = OpenSemaphore(EVENT_ALL_ACCESS, FALSE, L"mySemaphore");
HANDLE hMapFile = OpenFileMapping(FILE_MAP_ALL_ACCESS, NULL, L"ShareMemory");
LPVOID lpBase = MapViewOfFile(hMapFile, FILE_MAP_ALL_ACCESS, 0, 0, sizeof(info));

unordered_set<int> heapSet;  // 用于检查堆是否正确释放

/*
// 查看当前进程有哪些堆的句柄，调试使用
void ShowHeapHANDLE() {
    HANDLE hProcessHeap = GetProcessHeap();
    printf("hProcessHeap = %08X\n", (int)hProcessHeap);
    HANDLE hProcessHeaps[10] = { 0 };
    DWORD HandleNum = GetProcessHeaps(10, hProcessHeaps);
    for (int i = 0; i < (int)HandleNum; i++) {
        printf("hProcessHeaps[%d] = %08X\n", i, (int)hProcessHeaps[i]);
    }
    return;
}
*/

// 消息框操作
static int (WINAPI* OldMessageBoxW)(_In_opt_ HWND hWnd, _In_opt_ LPCWSTR lpText, _In_opt_ LPCWSTR lpCaption, _In_ UINT uType) = MessageBoxW;

static int (WINAPI* OldMessageBoxA)(_In_opt_ HWND hWnd, _In_opt_ LPCSTR lpText, _In_opt_ LPCSTR lpCaption, _In_ UINT uType) = MessageBoxA;

extern "C" __declspec(dllexport) int WINAPI NewMessageBoxA(_In_opt_ HWND hWnd, _In_opt_ LPCSTR lpText, _In_opt_ LPCSTR lpCaption, _In_ UINT uType) {
    // 向sendInfo中写入数据
    // 函数类型、系统时间以及参数个数
    sendInfo.type = MESSAGEBOXA;
    GetLocalTime(&(sendInfo.st));
    sendInfo.argNum = 4;
    // 参数名
    sprintf_s(sendInfo.argName[0], "hWnd");
    sprintf_s(sendInfo.argName[1], "lpText");
    sprintf_s(sendInfo.argName[2], "lpCaption");
    sprintf_s(sendInfo.argName[3], "uType");
    // 参数值
    sprintf_s(sendInfo.argValue[0], "%08X", hWnd);
    sprintf_s(sendInfo.argValue[1], "%s", lpText);
    sprintf_s(sendInfo.argValue[2], "%s", lpCaption);
    sprintf_s(sendInfo.argValue[3], "%08X", uType);

    
    // 将sendinfo赋值到共享内存
    memcpy(lpBase, &sendInfo, sizeof(info));
    // 进行V操作，使得信号量+1
    ReleaseSemaphore(hSemaphore, 1, NULL);
    sendInfo.argNum = 0;
    // 返回原始接口
    return OldMessageBoxA(hWnd, "Hooked!", "MessageBoxA", MB_OK);
}

extern "C" __declspec(dllexport) int WINAPI NewMessageBoxW(_In_opt_ HWND hWnd, _In_opt_ LPCWSTR lpText, _In_opt_ LPCWSTR lpCaption, _In_ UINT uType) {
    // 向sendInfo中写入数据
    // 函数类型、系统时间以及参数个数
    sendInfo.type = MESSAGEBOXW;
    GetLocalTime(&(sendInfo.st));
    sendInfo.argNum = 4;

    // 对lpText和lpCaption两个参数宽字节转char
    char lpTextChar[70];
    char lpCaptionChar[70];
    memset(lpTextChar, 0, sizeof(lpTextChar));
    memset(lpCaptionChar, 0, sizeof(lpCaptionChar));
    if (lpText != NULL) {
        WideCharToMultiByte(CP_ACP, 0, lpText, wcslen(lpText), lpTextChar, sizeof(lpTextChar), NULL, NULL);
    }
    if (lpCaption != NULL) {
        WideCharToMultiByte(CP_ACP, 0, lpCaption, wcslen(lpCaption), lpCaptionChar, sizeof(lpCaptionChar), NULL, NULL);
    }

    // 参数名
    sprintf_s(sendInfo.argName[0], "hWnd");
    sprintf_s(sendInfo.argName[1], "lpText");
    sprintf_s(sendInfo.argName[2], "lpCaption");
    sprintf_s(sendInfo.argName[3], "uType");
    // 参数值
    sprintf_s(sendInfo.argValue[0], "%08X", hWnd);
    strcpy_s(sendInfo.argValue[1], lpTextChar);
    strcpy_s(sendInfo.argValue[2], lpCaptionChar);
    sprintf_s(sendInfo.argValue[3], "%08X", uType);

    //-------------------------------------------------------------------- 与前端交互需要修改的地方
    memcpy(lpBase, &sendInfo, sizeof(info));
    ReleaseSemaphore(hSemaphore, 1, NULL);
    //--------------------------------------------------------------------

    // 复原argNum
    sendInfo.argNum = 0;
    // 返回原始接口
    return OldMessageBoxW(hWnd, L"Hooked!", L"MessageBoxW", MB_OK);
}

// 堆操作-----------------------------------------------------------------------------------------------------------------
static HANDLE(WINAPI* OldHeapCreate)(DWORD fIOoptions, SIZE_T dwInitialSize, SIZE_T dwMaximumSize) = HeapCreate;

static BOOL(WINAPI* OldHeapDestory)(HANDLE) = HeapDestroy;

static BOOL(WINAPI* OldHeapFree)(HANDLE hHeap, DWORD dwFlags, _Frees_ptr_opt_ LPVOID lpMem) = HeapFree;

extern "C" __declspec(dllexport)HANDLE WINAPI NewHeapCreate(DWORD fIOoptions, SIZE_T dwInitialSize, SIZE_T dwMaximumSize) {
    // 得到创建的堆的句柄
    HANDLE hFile = OldHeapCreate(fIOoptions, dwInitialSize, dwMaximumSize);

    // 向sendInfo中写入数据
    // 函数类型、系统时间以及参数个数
    sendInfo.type = HEAPCREATE;
    GetLocalTime(&(sendInfo.st));
    sendInfo.argNum = 4;

    // 参数名
    sprintf_s(sendInfo.argName[0], "fIOoptions");
    sprintf_s(sendInfo.argName[1], "dwInitialSize");
    sprintf_s(sendInfo.argName[2], "dwMaximumSize");
    sprintf_s(sendInfo.argName[3], "HANDLE");
    // 参数值
    sprintf_s(sendInfo.argValue[0], "%08X", fIOoptions);
    sprintf_s(sendInfo.argValue[1], "%08X", dwInitialSize);
    sprintf_s(sendInfo.argValue[2], "%08X", dwMaximumSize);
    sprintf_s(sendInfo.argValue[3], "%08X", hFile);

    // 记录已创建的堆的句柄
    unsigned temp = strtoul(sendInfo.argValue[3], NULL, 16);
    heapSet.insert(temp);

    //-------------------------------------------------------------------- 与前端交互需要修改的地方
    memcpy(lpBase, &sendInfo, sizeof(info));
    ReleaseSemaphore(hSemaphore, 1, NULL);
    //--------------------------------------------------------------------

    // 复原argNum
    sendInfo.argNum = 0;
    // 返回原始接口
    return OldHeapCreate(fIOoptions, dwInitialSize, dwMaximumSize);
}

extern "C" __declspec(dllexport) BOOL WINAPI NewHeapDestory(HANDLE hHeap) {
    // 向sendInfo中写入数据
    // 函数类型、系统时间以及参数个数
    sendInfo.type = HEAPCREATE;
    GetLocalTime(&(sendInfo.st));
    sendInfo.argNum = 1;

    // 参数名
    sprintf_s(sendInfo.argName[0], "hHeap");
    // 参数值
    sprintf_s(sendInfo.argValue[0], "%08X", hHeap);

    //-------------------------------------------------------------------- 与前端交互需要修改的地方
    memcpy(lpBase, &sendInfo, sizeof(info));
    ReleaseSemaphore(hSemaphore, 1, NULL);
    //--------------------------------------------------------------------

    // 复原argNum
    sendInfo.argNum = 0;
    // 返回原始接口
    return OldHeapDestory(hHeap);
}

extern "C" __declspec(dllexport) BOOL WINAPI NewHeapFree(HANDLE hHeap, DWORD dwFlags, _Frees_ptr_opt_ LPVOID lpMem) {
    // 向sendInfo中写入数据
    // 函数类型、系统时间以及参数个数
    sendInfo.type = HEAPFREE;
    GetLocalTime(&(sendInfo.st));
    sendInfo.argNum = 3;

    // 参数名
    sprintf_s(sendInfo.argName[0], "hHeap");
    sprintf_s(sendInfo.argName[1], "dwFlags");
    sprintf_s(sendInfo.argName[2], "lpMem");
    // 参数值
    sprintf_s(sendInfo.argValue[0], "%08X", hHeap);
    sprintf_s(sendInfo.argValue[1], "%08X", dwFlags);
    sprintf_s(sendInfo.argValue[2], "%08X", lpMem);

    // 判断是否是正常的释放堆
    unsigned temp = strtoul(sendInfo.argValue[0], NULL, 16);
    if (heapSet.find(temp) == heapSet.end()) {  // 释放了不存在的堆句柄
        sendInfo.argValue[9][9] = '1';          // 将 sendInfo.argValue[9][9] 的值改为了字符 '1'
    }
    else {
        heapSet.erase(temp);                    // 正常释放
    }

    //-------------------------------------------------------------------- 与前端交互需要修改的地方
    memcpy(lpBase, &sendInfo, sizeof(info));
    ReleaseSemaphore(hSemaphore, 1, NULL);
    //--------------------------------------------------------------------

    // 复原argNum和argValue
    sendInfo.argNum = 0;
    sendInfo.argValue[9][9] = '0';
    // 返回原始接口
    return OldHeapFree(hHeap, dwFlags, lpMem);
}

//文件操作
extern HANDLE(WINAPI* OldCreateFile)(
    LPCWSTR               lpFileName,
    DWORD                 dwDesiredAccess,
    DWORD                 dwShareMode,
    LPSECURITY_ATTRIBUTES lpSecurityAttributes,
    DWORD                 dwCreationDisposition,
    DWORD                 dwFlagsAndAttributes,
    HANDLE                hTemplateFile
    ) = CreateFile;

extern BOOL(WINAPI* OldWriteFile)(
    HANDLE       hFile,
    LPCVOID      lpBuffer,
    DWORD        nNumberOfBytesToWrite,
    LPDWORD      lpNumberOfBytesWritten,
    LPOVERLAPPED lpOverlapped
    ) = WriteFile;

extern BOOL(WINAPI* OldReadFile)(
    HANDLE       hFile,
    LPVOID       lpBuffer,
    DWORD        nNumberOfBytesToRead,
    LPDWORD      lpNumberOfBytesRead,
    LPOVERLAPPED lpOverlapped
    ) = ReadFile;

extern "C" __declspec(dllexport)HANDLE WINAPI NewCreateFile(
    LPCWSTR               lpFileName,
    DWORD                 dwDesiredAccess,
    DWORD                 dwShareMode,
    LPSECURITY_ATTRIBUTES lpSecurityAttributes,
    DWORD                 dwCreationDisposition,
    DWORD                 dwFlagsAndAttributes,
    HANDLE                hTemplateFile
);

extern "C" __declspec(dllexport)BOOL WINAPI NewWriteFile(
    HANDLE       hFile,
    LPCVOID      lpBuffer,
    DWORD        nNumberOfBytesToWrite,
    LPDWORD      lpNumberOfBytesWritten,
    LPOVERLAPPED lpOverlapped
);

extern "C" __declspec(dllexport)BOOL WINAPI NewReadFile(
    HANDLE       hFile,
    LPVOID       lpBuffer,
    DWORD        nNumberOfBytesToRead,
    LPDWORD      lpNumberOfBytesRead,
    LPOVERLAPPED lpOverlapped
);

//注册表操作
static LSTATUS(WINAPI* OldRegCreateKeyEx)(
    HKEY                        hKey,
    LPCWSTR                     lpSubKey,
    DWORD                       Reserved,
    LPWSTR                      lpClass,
    DWORD                       dwOptions,
    REGSAM                      samDesired,
    const LPSECURITY_ATTRIBUTES lpSecurityAttributes,
    PHKEY                       phkResult,
    LPDWORD                     lpdwDisposition) = RegCreateKeyEx;

static LSTATUS(WINAPI* OldRegSetValueEx)(
    HKEY       hKey,
    LPCWSTR    lpValueName,
    DWORD      Reserved,
    DWORD      dwType,
    const BYTE* lpData,
    DWORD      cbData
    ) = RegSetValueEx;

static LSTATUS(WINAPI* OldRegCloseKey)(HKEY hKey) = RegCloseKey;

static LSTATUS(WINAPI* OldRegOpenKeyEx)(
    HKEY    hKey,
    LPCWSTR lpSubKey,
    DWORD   ulOptions,
    REGSAM  samDesired,
    PHKEY   phkResult
    ) = RegOpenKeyEx;

static LSTATUS(WINAPI* OldRegDeleteValue)(
    HKEY    hKey,
    LPCWSTR lpValueName
    ) = RegDeleteValue;

extern "C" __declspec(dllexport)LSTATUS WINAPI NewRegCreateKeyEx(
    HKEY                        hKey,
    LPCWSTR                     lpSubKey,
    DWORD                       Reserved,
    LPWSTR                      lpClass,
    DWORD                       dwOptions,
    REGSAM                      samDesired,
    const LPSECURITY_ATTRIBUTES lpSecurityAttributes,
    PHKEY                       phkResult,
    LPDWORD                     lpdwDisposition
) {
    char temp[70];
    sendInfo.argNum = 9;
    // 参数名
    sprintf_s(sendInfo.argName[0], "hKey");
    sprintf_s(sendInfo.argName[1], "lpSubKey");
    sprintf_s(sendInfo.argName[2], "Reserved");
    sprintf_s(sendInfo.argName[3], "lpClass");
    sprintf_s(sendInfo.argName[4], "dwOptions");
    sprintf_s(sendInfo.argName[5], "samDesired");
    sprintf_s(sendInfo.argName[6], "lpSecurityAttributes");
    sprintf_s(sendInfo.argName[7], "phkResult");
    sprintf_s(sendInfo.argName[8], "lpdwDisposition");
    // 参数值
    sprintf_s(sendInfo.argValue[0], "%08X", hKey);
    // 宽字节转char
    memset(temp, 0, sizeof(temp));
    WideCharToMultiByte(CP_ACP, 0, lpSubKey, wcslen(lpSubKey), temp, sizeof(temp), NULL, NULL);
    strcpy_s(sendInfo.argValue[1], temp);
    sprintf_s(sendInfo.argValue[2], "%08X", Reserved);
    sprintf_s(sendInfo.argValue[3], "%08X", lpClass);
    sprintf_s(sendInfo.argValue[4], "%08X", dwOptions);
    sprintf_s(sendInfo.argValue[5], "%08X", samDesired);
    sprintf_s(sendInfo.argValue[6], "%08X", lpSecurityAttributes);
    sprintf_s(sendInfo.argValue[7], "%08X", phkResult);
    sprintf_s(sendInfo.argValue[8], "%08X", lpdwDisposition);

    sendInfo.type = REGCREATEKEYEX;
    GetLocalTime(&(sendInfo.st));
    memcpy(lpBase, &sendInfo, sizeof(info));
    ReleaseSemaphore(hSemaphore, 1, NULL);
    return OldRegCreateKeyEx(hKey, lpSubKey, Reserved, lpClass, dwOptions, samDesired, lpSecurityAttributes, phkResult, lpdwDisposition);
}

extern "C" __declspec(dllexport)LSTATUS WINAPI NewRegSetValueEx(
    HKEY       hKey,
    LPCWSTR    lpValueName,
    DWORD      Reserved,
    DWORD      dwType,
    const BYTE * lpData,
    DWORD      cbData)
{
    char temp[70];
    sendInfo.argNum = 6;
    // 参数名
    sprintf_s(sendInfo.argName[0], "hKey");
    sprintf_s(sendInfo.argName[1], "lpValueName");
    sprintf_s(sendInfo.argName[2], "Reserved");
    sprintf_s(sendInfo.argName[3], "dwType");
    sprintf_s(sendInfo.argName[4], "lpData");
    sprintf_s(sendInfo.argName[5], "cbData");
    // 参数值
    sprintf_s(sendInfo.argValue[0], "%08X", hKey);
    // 宽字节转char
    memset(temp, 0, sizeof(temp));
    WideCharToMultiByte(CP_ACP, 0, lpValueName, wcslen(lpValueName), temp, sizeof(temp), NULL, NULL);
    strcpy_s(sendInfo.argValue[1], temp);
    sprintf_s(sendInfo.argValue[2], "%08X", Reserved);
    sprintf_s(sendInfo.argValue[3], "%08X", dwType);
    memset(temp, 0, sizeof(temp));
    WideCharToMultiByte(CP_ACP, 0, (LPCWCH)lpData, wcslen((LPCWCH)lpData), temp, sizeof(temp), NULL, NULL);
    strcpy_s(sendInfo.argValue[4], temp);
    //strcpy(sendInfo.argValue[4], (const char *)lpData);
    //sprintf_s(sendInfo.argValue[4], "%08X", lpData);
    sprintf_s(sendInfo.argValue[5], "%08X", cbData);


    sendInfo.type = REGSETVALUEEX;
    GetLocalTime(&(sendInfo.st));
    memcpy(lpBase, &sendInfo, sizeof(info));
    ReleaseSemaphore(hSemaphore, 1, NULL);
    return OldRegSetValueEx(hKey, lpValueName, Reserved, dwType, lpData, cbData);
}

extern "C" __declspec(dllexport)LSTATUS WINAPI NewRegCloseKey(HKEY hKey)
{
    sendInfo.argNum = 1;
    // 参数名
    sprintf_s(sendInfo.argName[0], "hKey");
    // 参数值
    sprintf_s(sendInfo.argValue[0], "%08X", hKey);
    sendInfo.type = REGCLOSEKEY;
    GetLocalTime(&(sendInfo.st));
    memcpy(lpBase, &sendInfo, sizeof(info));
    ReleaseSemaphore(hSemaphore, 1, NULL);
    return OldRegCloseKey(hKey);
}

extern "C" __declspec(dllexport)LSTATUS WINAPI NewRegOpenKeyEx(
    HKEY    hKey,
    LPCWSTR lpSubKey,
    DWORD   ulOptions,
    REGSAM  samDesired,
    PHKEY   phkResult)
{
    char temp[70];
    sendInfo.argNum = 5;
    // 参数名
    sprintf_s(sendInfo.argName[0], "hKey");
    sprintf_s(sendInfo.argName[1], "lpSubKey");
    sprintf_s(sendInfo.argName[2], "ulOptions");
    sprintf_s(sendInfo.argName[3], "samDesired");
    sprintf_s(sendInfo.argName[4], "phkResult");
    // 参数值
    sprintf_s(sendInfo.argValue[0], "%08X", hKey);
    // 宽字节转char
    memset(temp, 0, sizeof(temp));
    WideCharToMultiByte(CP_ACP, 0, lpSubKey, wcslen(lpSubKey), temp, sizeof(temp), NULL, NULL);
    strcpy_s(sendInfo.argValue[1], temp);
    sprintf_s(sendInfo.argValue[2], "%08X", ulOptions);
    sprintf_s(sendInfo.argValue[3], "%08X", samDesired);
    sprintf_s(sendInfo.argValue[4], "%08X", phkResult);

    sendInfo.type = REGOPENKEYEX;
    GetLocalTime(&(sendInfo.st));
    memcpy(lpBase, &sendInfo, sizeof(info));
    ReleaseSemaphore(hSemaphore, 1, NULL);
    return OldRegOpenKeyEx(hKey, lpSubKey, ulOptions, samDesired, phkResult);
}

extern "C" __declspec(dllexport)LSTATUS WINAPI NewRegDeleteValue(
    HKEY    hKey,
    LPCWSTR lpValueName)
{
    char temp[70];
    sendInfo.argNum = 2;
    // 参数名
    sprintf_s(sendInfo.argName[0], "hKey");
    sprintf_s(sendInfo.argName[1], "lpValueName");
    // 参数值
    sprintf_s(sendInfo.argValue[0], "%08X", hKey);
    // 宽字节转char
    memset(temp, 0, sizeof(temp));
    WideCharToMultiByte(CP_ACP, 0, lpValueName, wcslen(lpValueName), temp, sizeof(temp), NULL, NULL);
    strcpy_s(sendInfo.argValue[1], temp);
    sendInfo.type = REGDELETEVALUE;
    GetLocalTime(&(sendInfo.st));
    memcpy(lpBase, &sendInfo, sizeof(info));
    ReleaseSemaphore(hSemaphore, 1, NULL);
    return OldRegDeleteValue(hKey, lpValueName);
}

BOOL WINAPI DllMain(HMODULE hModule,
    DWORD ul_reason_for_call,
    LPVOID lpReserved
)
{

    switch (ul_reason_for_call)
    {
    case DLL_PROCESS_ATTACH:
    {
        DisableThreadLibraryCalls(hModule);
        DetourTransactionBegin();
        DetourUpdateThread(GetCurrentThread());
        DetourAttach(&(PVOID&)OldMessageBoxW, NewMessageBoxW);
        DetourAttach(&(PVOID&)OldMessageBoxA, NewMessageBoxA);
        DetourAttach(&(PVOID&)OldCreateFile, NewCreateFile);
        DetourAttach(&(PVOID&)OldWriteFile, NewWriteFile);
        DetourAttach(&(PVOID&)OldReadFile, NewReadFile);
        DetourAttach(&(PVOID&)OldHeapCreate, NewHeapCreate);
        DetourAttach(&(PVOID&)OldHeapDestory, NewHeapDestory);
        DetourAttach(&(PVOID&)OldHeapFree, NewHeapFree);
        DetourAttach(&(PVOID&)OldRegCreateKeyEx, NewRegCreateKeyEx);
        DetourAttach(&(PVOID&)OldRegSetValueEx, NewRegSetValueEx);
        DetourAttach(&(PVOID&)OldRegDeleteValue, NewRegDeleteValue);
        DetourAttach(&(PVOID&)OldRegCloseKey, NewRegCloseKey);
        DetourAttach(&(PVOID&)OldRegOpenKeyEx, NewRegOpenKeyEx);
        DetourTransactionCommit();
        break;
    }
    case DLL_THREAD_ATTACH:
    case DLL_THREAD_DETACH:
    case DLL_PROCESS_DETACH:
    {
        DetourTransactionBegin();
        DetourUpdateThread(GetCurrentThread());
        DetourDetach(&(PVOID&)OldMessageBoxW, NewMessageBoxW);
        DetourDetach(&(PVOID&)OldMessageBoxA, NewMessageBoxA);
        DetourDetach(&(PVOID&)OldCreateFile, NewCreateFile);
        DetourDetach(&(PVOID&)OldWriteFile, NewWriteFile);
        DetourDetach(&(PVOID&)OldReadFile, NewReadFile);
        DetourDetach(&(PVOID&)OldHeapCreate, NewHeapCreate);
        DetourDetach(&(PVOID&)OldHeapDestory, NewHeapDestory);
        DetourDetach(&(PVOID&)OldHeapFree, NewHeapFree);
        DetourDetach(&(PVOID&)OldRegCreateKeyEx, NewRegCreateKeyEx);
        DetourDetach(&(PVOID&)OldRegSetValueEx, NewRegSetValueEx);
        DetourDetach(&(PVOID&)OldRegDeleteValue, NewRegDeleteValue);
        DetourDetach(&(PVOID&)OldRegCloseKey, NewRegCloseKey);
        DetourDetach(&(PVOID&)OldRegOpenKeyEx, NewRegOpenKeyEx);
        DetourTransactionCommit();
        UnmapViewOfFile(lpBase);
        CloseHandle(hMapFile);
        break;
    }
    }
    return true;
}